package SimpleLACore

import chisel3._
import chisel3.util._
import chisel3.util.random.LFSR

class SimpleICacheAXI(val id: Int) extends Module with SimpleLACoreParam {
  val io = IO(new Bundle {
    val core = new SimpleICacheIO
    val mem = new AXI4Bundle
  })
  require((nSets & (nSets - 1)) == 0, "nSets must be power of 2")
  require((lineWords & (lineWords - 1)) == 0, "wordsPerLine must be power of 2")
  require((nWays & (nWays - 1)) == 0, "nWays must be power of 2")

  def getIdx(addr: UInt):UInt = addr(idxBits + lgLineBytes - 1, lgLineBytes)
  def getTag(addr: UInt):UInt = addr(31, 32 - tagBits)
  def getOffset(addr: UInt):UInt = addr(lgLineBytes - 1, 2)

  val rand = LFSR(8)(lgnWays - 1, 0)

  val reset_counter = RegInit(0.U((idxBits + 1).W))
  val reset_done = reset_counter(idxBits)
  val metaArray = Seq.fill(nWays)(Mem(nSets, new SimpleCacheMeta))
  val dataArray = Seq.fill(nWays)(Mem(nSets, Vec(lineWords, Vec(4, UInt(8.W)))))
  //rw ports for Arrays
  val index = WireInit(UInt(idxBits.W), DontCare)
  val metaPort = metaArray.map(_(index))
  val dataPort = dataArray.map(_(index))

  val coreBusyReg = RegInit(false.B)
  val coreReqReg = RegEnable(io.core.req.bits, io.core.req.valid)
  val coreBusy = coreBusyReg || io.core.req.valid
  val coreReq = Mux(io.core.req.valid, io.core.req.bits, coreReqReg)
  val coreMiss = WireInit(false.B)

  val buf = Reg(Vec(lineWords, Vec(4, UInt(8.W))))
  val bufBase = Reg(UInt(lgLineWords.W))
  val bufAddr = Reg(UInt((32 - lgLineBytes).W))
  val bufHead = Reg(Vec(lineWords, Bool()))
  val bufHeadNext = (bufHead.reverse.tail :+ bufHead.last).reverse
  val bufValid = RegInit(VecInit.tabulate(lineWords)(_=>false.B))
  val bufInvalid = RegInit(false.B)

  val sReady :: sFill :: Nil = Enum(2)
  val state = RegInit(sReady)
  val aFire = RegInit(false.B)

  io.mem.ar.valid := false.B
  io.mem.aw.valid := false.B
  io.mem.w.valid  := false.B
  io.mem.ar.bits := DontCare
  io.mem.aw.bits := DontCare
  io.mem.w.bits  := DontCare
  io.mem.ar.bits.id := id.U
  io.mem.ar.bits.size := 2.U
  io.mem.ar.bits.lock := 0.U
  io.mem.ar.bits.cache := 0.U
  io.mem.ar.bits.prot := 0.U
  io.mem.r.ready  := true.B
  io.mem.b.ready  := true.B
  io.core.resp.valid := false.B
  io.core.resp.bits := DontCare

  when(io.core.resp.valid) {
    coreBusyReg := false.B
  }.elsewhen(io.core.req.valid) {
    coreBusyReg := true.B
  }
  //reset
  when(!reset_done){
    reset_counter := reset_counter + 1.U
    metaArray.foreach(_(reset_counter.tail(1)).valid := false.B)
  }
  //cacop
  when(reset_done && io.core.cacop.valid){
    index := getIdx(io.core.cacop.bits.addr)
    val waySel = UIntToOH(io.core.cacop.bits.addr(lgnWays - 1, 0)).asBools
    when(io.core.cacop.bits.code =/= 2.U){
      (metaPort zip waySel).foreach{ case (meta, sel) => when(sel){meta.valid := false.B}}
    }.otherwise{
      when(bufAddr === io.core.cacop.bits.addr(31, lgLineBytes)){bufInvalid := true.B}
      metaPort.foreach{ meta => when(meta.tag === getTag(io.core.cacop.bits.addr)) {meta.valid := false.B}}
    }
  }
  //req
  when(reset_done && coreBusy){
    when(coreReq.mat =/= 0.U){//CC
      //Use another Port for req
      val metas = metaArray.map(_.read(getIdx(coreReq.addr)))
      val datum = dataArray.map(_.read(getIdx(coreReq.addr)))
      //buffer hit?
      val bufHit = bufAddr === coreReq.addr(31, lgLineBytes) && bufValid(getOffset(coreReq.addr))
      val bufData = buf(getOffset(coreReq.addr)).asUInt & Fill(32, bufHit)
      //cache hit?
      val (hit, data) = (datum zip metas).foldLeft((bufHit, bufData)) { case ((hitIn, dataIn), (line, meta)) =>
        val hit = meta.valid && meta.tag === getTag(coreReq.addr)
        val data = line(getOffset(coreReq.addr))
        (hit | hitIn, data.asUInt & Fill(32, hit) | dataIn)
      }
      when(hit){
        io.core.resp.valid := true.B
        io.core.resp.bits := data
      }.otherwise{
        coreMiss := true.B
      }
    }.elsewhen(state === sReady){//SUC
      when(!aFire){
        io.mem.ar.valid := true.B
        io.mem.ar.bits.addr := coreReq.addr
        io.mem.ar.bits.len := 0.U
        io.mem.ar.bits.burst := 1.U
      }
      when(io.mem.r.valid){
        aFire := false.B
        io.core.resp.valid := true.B
        io.core.resp.bits := io.mem.r.bits.data
      }.elsewhen(io.mem.ar.fire){
        aFire := true.B
      }
    }
  }

  when(reset_done){
    switch(state) {
      is(sReady){
        when(coreMiss){
          state := sFill
          bufAddr := coreReq.addr(31, lgLineBytes)
          bufBase := getOffset(coreReq.addr)
          bufHead := UIntToOH(getOffset(coreReq.addr)).asBools
          bufInvalid := false.B
        }
      }
      is(sFill) {
        when(!aFire){
          io.mem.ar.valid := true.B
          io.mem.ar.bits.addr := Cat(bufAddr, bufBase, 0.U(2.W))
          io.mem.ar.bits.len := (lineWords - 1).U
          io.mem.ar.bits.burst := 2.U
          when(io.mem.ar.ready){
            aFire := true.B
          }
        }
        when(io.mem.r.valid) {
          (buf zip bufValid zip bufHead).foreach{ case ((word, valid), isHead) =>
            when(isHead){
              word := io.mem.r.bits.data.asTypeOf(word)
              valid := true.B
            }
          }
          bufHead := bufHeadNext
        }
        //refill
        when(bufValid.reduceTree(_&_) && !io.core.cacop.valid){//Avoid write simultaneously && Cacop may invalid buffer
          when(!bufInvalid){
            val replaceWay = rand // random replacement
            index := bufAddr // scala will truncate tag bits
            val waySel = UIntToOH(replaceWay).asBools
            //write data and meta
            (dataPort zip waySel).foreach{ case (line, sel) =>
              when(sel){
                line := buf
              }
            }
            (metaPort zip waySel).foreach{ case (meta, sel) =>
              when(sel){
                meta.tag := bufAddr(tagBits + idxBits - 1, idxBits)
                meta.valid := true.B
              }
            }
          }
          aFire := false.B
          state := sReady
          bufValid.foreach(_ := false.B)
        }
      }
    }
  }
}